#include "demux.h"

DemuxServer::DemuxServer(EventLoop *loop,
                         const InetAddress &listenAddr,
                         const InetAddress &socksAddress)
    : _loop(loop),
      _server(loop, listenAddr),
      _sockAddress(socksAddress)
{
    _server.setConnectionCallback(std::bind(&DemuxServer::onServerConnection, this, _1));
    _server.setMessageCallback(std::bind(&DemuxServer::onServerMessage, this, _1, _2, _3));
}

void DemuxServer::start()
{
    _server.start();
}

void DemuxServer::onServerConnection(const TcpConnectionPtr &conn)
{
    if (conn->connected())
    {
        if (_serverConn)
        {
            conn->shutdown();
        }
        else
        {
            _serverConn = conn;
            LOG_INFO << "onServerConnection set _serverConn";
        }
    }
    else
    {
        if (conn == _serverConn)
        {
            _serverConn.reset();
            _sockConnections.clear();
            LOG_INFO << "onServerConnection reset _serverConn";
        }
    }
}
void DemuxServer::onServerMessage(const TcpConnectionPtr &conn, Buffer *buffer, Timestamp)
{
    LOG_DEBUG;
    while (buffer->readableBytes() > kHeaderLen)
    {
        int len_message = static_cast<uint8_t>(*buffer->peek());
        if (buffer->readableBytes() < kHeaderLen + len_message)
        {
            break;
        }
        else
        {
            int connId = static_cast<uint8_t>(buffer->peek()[1]);
            connId |= (static_cast<uint8_t>(buffer->peek()[2]) << 8);
            if (connId != 0)
            {
                assert(_sockConnections.find(connId) != _sockConnections.end());
                TcpConnectionPtr &sockConn = _sockConnections[connId].clientConn;
                if (sockConn)
                {
                    assert(_sockConnections[connId].pending.readableBytes() == 0);
                    sockConn->send(buffer->peek() + kHeaderLen, len_message);
                }
                else
                {
                    _sockConnections[connId].pending.append(buffer->peek() + kHeaderLen, len_message);
                }
            }
            else//建立主服务器的连接
            {
                string cmd(buffer->peek() + kHeaderLen, len_message);
                doCommand(cmd);
            }
        }
        buffer->retrieve(kHeaderLen + len_message);
    }
    LOG_DEBUG;
}
void DemuxServer::doCommand(const string &cmd)
{
    static const string kCmd = "CONN ";
    int connId = atoi(&cmd[kCmd.size()]);
    bool isUp = cmd.find(" IS UP") != std::string::npos;
    LOG_INFO << "doCommand " << connId << " " << isUp;
    if (isUp)
    {
        assert(_sockConnections.find(connId) == _sockConnections.end());
        char connName[256];
        snprintf(connName, sizeof(connName), "SocksClient %d", connId);

        client_conn node;
        node.connID = connId;
        node.client.reset(new TcpClient(_loop, _sockAddress));
        node.client->setConnectionCallback(std::bind(&DemuxServer::onSocksConnection, this, connId, _1));
        node.client->setMessageCallback(std::bind(&DemuxServer::onSocksMessage, this, connId, _1, _2, _3));
        _sockConnections[connId] = node;

        node.client->connect(); //连接
    }
    else
    {
        assert(_sockConnections.find(connId) != _sockConnections.end());
        TcpConnectionPtr &socksConn = _sockConnections[connId].clientConn;
        if (socksConn)
        {
            socksConn->shutdown();
        }
        else
        {
            _sockConnections.erase(connId);
        }
    }
}

void DemuxServer::onSocksConnection(int connId,
                                    const TcpConnectionPtr &conn)
{
    assert(_sockConnections.find(connId) != _sockConnections.end());
    if (conn->connected())
    {
        _sockConnections[connId].clientConn = conn;
        Buffer &pendingData = _sockConnections[connId].pending;
        if (pendingData.readableBytes() > 0)
        {
            conn->send(&pendingData); //发送之前缓存的buffer
        }
    }
    else
    {
        if (_serverConn)
        {
            char buffer[256];
            int len_message = snprintf(buffer, sizeof(buffer), "DISCONNECT %d\r\n", connId);
            Buffer packet;
            packet.append(buffer, len_message);
            sendServerPacket(0, &packet);
        }
        else
        {
            _sockConnections.erase(connId);
        }
    }
}
void DemuxServer::onSocksMessage(int connId,
                                 const TcpConnectionPtr &conn,
                                 Buffer *buffer, Timestamp)
{
    LOG_DEBUG;
    assert(_sockConnections.find(connId) != _sockConnections.end());
    while (buffer->readableBytes() > kMaxPacketLen)
    {
        Buffer packet;
        packet.append(buffer->peek(), kMaxPacketLen);
        buffer->retrieve(kMaxPacketLen);
        sendServerPacket(connId, &packet);
    }
    if (buffer->readableBytes() > 0)
    {
        sendServerPacket(connId, buffer);
    }
    LOG_DEBUG;
}

void DemuxServer::sendServerPacket(int connId, Buffer *buffer)
{
    size_t len_message = buffer->readableBytes();

    LOG_DEBUG << len_message;
    assert(len_message <= kMaxPacketLen);

    uint8_t header[kHeaderLen] =
    {
        static_cast<uint8_t>(len_message),
        static_cast<uint8_t>(connId & 0xFF),
        static_cast<uint8_t>((connId & 0xFF00) >> 8)
    };

    buffer->prepend(header,kHeaderLen);
    if(_serverConn)
    {
        _serverConn->send(buffer);
    }
}